/*
 * This demo application implements
 * some nice functions related to cryptography
 *
 * Author: Liviu Itoafa <liviu_itoafaATgmailDOTcom>
 *
 */

#include <limits.h>
#include <openssl/crypto.h>
#include <openssl/sha.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "functions.h"
#include "trace.h"
#include "util.h"

#define SUB_MSG_LEN             (20u)

// set global trace level
unsigned int traceLevel = TRACE_LEVEL_INFO;

static void print_help(char *prog_name);

int main(int argc, char **argv){
    char *prog_name = argv[0];


    if(argc<2) {
        TRACE_WARNING("Too few argumens\n");
        print_help(prog_name);
        return 1;
    }

    if(strcmp(argv[1],"--help") == 0){
        // print usage info
        print_help(prog_name);
        return 0;
    }
    else if(strcmp(argv[1],"--inv-mod") == 0){
        // compute inverse module for small numbers
        if(argc != 4) {
            TRACE_ERROR("Invalid number of arguments\n");
        }

        TRACE_THIS("inv(%s) mod %s = %d\n", argv[2], argv[3],
                inv_mod(atoi(argv[2]), atoi(argv[3])));
    }
    else if (strcmp(argv[1],"--inv-mod-bn") == 0){
        // compute inverse module for big numbers
        char *rez;
        inv_mod_bn(argv[2], argv[3], &rez);
        TRACE_THIS("inv(%s) mod %s = %s\n", argv[2], argv[3], rez);
        OPENSSL_free(rez);
    }
    else if (strcmp(argv[1],"--exp-mod-bn") == 0){
        // compute exponentiation modulo n
        char *rez_exp;

        exp_mod_bn(argv[2], argv[3], argv[4], &rez_exp);
        TRACE_THIS("%s^%s mod %s = %s\n", argv[2], argv[3], argv[4], rez_exp);
        OPENSSL_free(rez_exp);
    }
    else if (strcmp(argv[1],"--crt") == 0){
        // compute solution to chinese reminder theorem equations
        int count = atoi(argv[2]);
        if((count == 0) || (count == INT_MAX) || (count == INT_MIN)) {
            TRACE_ERROR("Invalid number of equations\n");
            return 1;
        }

        if(argc != (2*count+3)) {
            TRACE_ERROR("Invalid number of coefficients\n");
            return 1;
        }

        int *a = (int*) malloc(count * sizeof(int));
        int *m = (int*) malloc(count * sizeof(int));

        int bOK = 1;

        for(int i=0; i<count; i++) {
            a[i] = atoi(argv[3+i]);
            m[i] = atoi(argv[3+count+i]);
            if((a[i] == 0) || (a[i] == INT_MAX) || (a[i] == INT_MIN)) {
                TRACE_ERROR("Invalid coefficient. Could not decode %s\n", argv[3+i]);
                bOK = 0;
                break;
            }
            if((m[i] == 0) || (m[i] == INT_MAX) || (m[i] == INT_MIN)) {
                TRACE_ERROR("Invalid coefficient. Could not decode %s\n", argv[3+count+i]);
                bOK = 0;
                break;
            }
        }

        if(!bOK) { // one parameter decoded wrong. Break case label
            TRACE_ERROR("Parameter couldn't be decoded\n");
            return 1;
        }

        int rez = 0;
        int M = 0;

        int rc = crt(a, m, count, &rez, &M);
        if(0 == rc) {
            TRACE_THIS("crt(ai,mi)=%d mod %d\n", rez, M);
        }
        else {
            TRACE_ERROR("Error in crt function\n");
            return 1;
        }
    }
    else if (strcmp(argv[1],"--crt-bn") == 0){
        // compute solution to crt using BIGNUMs
        int count = atoi(argv[2]);
        if((count == 0) || (count == INT_MAX) || (count == INT_MIN)) {
            TRACE_ERROR("Invalid number of equations\n");
            return 1;
        }

        if(argc != (2*count+3)) {
            TRACE_ERROR("Invalid number of coefficients\n");
            return 1;
        }

        char *rez;
        char *M;

        int rc = crt_bn(argv+3, argv+3+count, count, &rez, &M);
        if(0 == rc) {
            TRACE_THIS("crt_bn(ai,mi)=%s mod %s\n", rez, M);
        }
        else {
            TRACE_ERROR("Error in crt_bn function\n");
            return 1;
        }

        OPENSSL_free(rez);
        OPENSSL_free(M);
    }

    else if (strcmp(argv[1],"--gcd") == 0){
        // compute greatest common divisor
        TRACE_THIS("gcd(%d,%d)=%d\n", atoi(argv[2]), atoi(argv[3]),
                gcd(atoi(argv[2]), atoi(argv[3])));
    }
    else if (strcmp(argv[1],"--enc-b64") == 0){
        // base64 encoding
        int len = strlen(argv[2]);

        char *out_buf = NULL;

        int out_len = 0;

        base64_enc((const unsigned char *)argv[2], len, &out_buf, &out_len);

        TRACE_DEBUG("Input string length: %d\n", len);
        TRACE_THIS("Encoded string: %s\n", out_buf);
        TRACE_DEBUG("Encoded string length: %d\n", out_len);

        free(out_buf);
    }
    else if (strcmp(argv[1],"--dec-b64") == 0){
        // base64 decoding
        int len = strlen(argv[2]);

        char *out_buf = NULL;

        int out_len = 0;

        base64_dec((unsigned char *)argv[2], len, &out_buf, &out_len);

        TRACE_DEBUG("Input string length: %d\n", len);
        TRACE_THIS("Decoded string: %s\n", out_buf);
        TRACE_DEBUG("Decoded string length: %d\n", out_len);

        free(out_buf);
    }
    else if (strcmp(argv[1],"--gen-keys") == 0){
        // generate private/public key pair
        if(3 == argc) {
            // passphrase was provided
            if(generate_DSA_key_pair(argv[2])) {
                TRACE_ERROR("generate_DSA_key_pair error\n");
                return 1;
            }
        }
        else {
            if( 2 == argc){
                // prompt for password
                if(generate_DSA_key_pair("")) {
                    TRACE_ERROR("generate_DSA_key_pair error\n");
                    return 1;
                }
            }
            else {
                TRACE_ERROR("Invalid number of arguments\n");
                return 1;
            }
        }
    }
    else if (strcmp(argv[1],"--sha1") == 0){
        // compute sha1 hash of a file
        unsigned char *hash = (unsigned char *)malloc(SHA_DIGEST_LENGTH);
        if(NULL == hash) {
            TRACE_ERROR("malloc error\n");
            return 1;
        }

        unsigned int hash_len;

        if(calc_sha1(argv[2], &hash, &hash_len)){
            TRACE_ERROR("calc_sha1 error\n");
            return 1;
        }

        TRACE_THIS("Digest is: ");
        for(unsigned int i=0; i<hash_len; i++){
            TRACE_THIS("%02x", hash[i]);
        }
        TRACE_THIS(" hash length:%d\n", hash_len);
    }
    else if (strcmp(argv[1],"--sign") == 0){
        // sign file with private key from pem file
        if(argc<4) {
            print_help(prog_name);
            return 1;
        }
        if(4 == argc){
            //ask for pem file password
            if(dsa_sign_file(argv[2], argv[3], NULL, NULL, 0, 0)){
                TRACE_ERROR("dsa_sign_file failed\n");
                return 1;
            }
        }
        else{
            //use passphrase to open private key
            if(dsa_sign_file(argv[2], argv[3], argv[4], NULL, 0, 0)){
                TRACE_ERROR("dsa_sign_file failed\n");
                return 1;
            }
        }
    }
    else if (strcmp(argv[1],"--sign-with-leak") == 0){
        // sign file with private key from pem file
        if(argc<4) {
            print_help(prog_name);
            return 1;
        }
        if(4 == argc){
            //ask for pem file password
            if(dsa_sign_file(argv[2], argv[3], NULL, NULL, 0, 1)){
                TRACE_ERROR("dsa_sign_file failed\n");
                return 1;
            }
        }
        else{
            //use passphrase to open private key
            if(dsa_sign_file(argv[2], argv[3], argv[4], NULL, 0, 1)){
                TRACE_ERROR("dsa_sign_file failed\n");
                return 1;
            }
        }
    }
    else if (strcmp(argv[1], "--sign-with-msg")==0){
        if(argc<4) {
            print_help(prog_name);
            return 1;
        }
        if((argc == 4) || (argc ==5)){
            //ask for subliminal message and for pem file password
            char *sub_msg;

            sub_msg = (char *) malloc (SUB_MSG_LEN);
            memset(sub_msg, SUB_MSG_LEN, 0);

            printf("Enter the hidden message (20 chars max):\n");
            fgets(sub_msg, SUB_MSG_LEN, stdin);

            sub_msg[strlen(sub_msg)-1]='\0'; // replace newline with '\0'

            if(4 == argc){
                if(dsa_sign_file(argv[2], argv[3], NULL, sub_msg, strlen(sub_msg), 0)){
                    TRACE_ERROR("dsa_sign_file failed\n");
                    return 1;
                }
            }
            else { //argc == 5
                if(dsa_sign_file(argv[2], argv[3], argv[4], sub_msg, strlen(sub_msg), 0)){
                    TRACE_ERROR("dsa_sign_file failed\n");
                    return 1;
                }
            }
        }
        else {
            if(dsa_sign_file(argv[2], argv[3], argv[4], argv[5], strlen(argv[5]),0)){
                TRACE_ERROR("dsa_sign_file failed\n");
                return 1;
            }
        }
    }
    else if (strcmp(argv[1], "--verify") == 0){
        // verify signature of file
        if(verif_sig_file(argv[2], argv[3], argv[4])){
            TRACE_ERROR("Error verifying signature\n");
            return 1;
        }
        else {
            TRACE_THIS("Signature OK\n");
        }
    }
    else if (strcmp(argv[1], "--get-msg") == 0){
        //get message embedded into k

        unsigned char *msg = NULL;

        if(5 == argc){
            //password not provided. Ask for one
            if(dsa_get_k_with_x(argv[2], argv[3], argv[4], NULL, &msg)){
                TRACE_ERROR("Error getting subliminal message\n");
                return 1;
            }
        }
        else {
            // password given
            if(dsa_get_k_with_x(argv[2], argv[3], argv[4], argv[5], &msg)){
                TRACE_ERROR("Error getting subliminal message\n");
                return 1;
            }
        }
        TRACE_THIS("Got: %s\n", msg);

        free(msg);
    }
    else if (strcmp(argv[1], "--is-quad-res") == 0) {
        // test for quadratic residue
        if(argc != 4) {
            TRACE_ERROR("Invalid number of parameters\n");
            return 1;
        }

        int ret = is_quad_res(argv[2], argv[3]);

        TRACE_THIS("%s is quadratic %sresidue modulo %s\n", argv[2],
                ret==1?"":"non", argv[3]);
    }
    else {
        TRACE_WARNING("Bad option %s\n\n", argv[1]);
        print_help(prog_name);
        return 1;
    }

    return 0;
}

static void print_help(char *prog_name) {
    TRACE_THIS("Usage: %s  -OP <param 1> ... <param N>\n\n", prog_name);
    TRACE_THIS("Use %s --help for help on functions usage\n\n", prog_name);
    TRACE_THIS("Arithmetic functions\n");
    TRACE_THIS(" --inv-mod [nr] [mod]\n\tinverse modulo\n");
    TRACE_THIS(" --inv-mod-bn [nr] [mod]\n\tinverse modulo (big numbers)\n");
    TRACE_THIS(" --exp-mod-bn [nr] [pow] [mod]\n\ta^p mod m (big numbers)\n");
    TRACE_THIS(" --gcd [a] [b]\n\tgcd(a,b)\n");
    TRACE_THIS("\nCrypto functions\n");
    TRACE_THIS(" --crt [k] [a1] .. [ak] [m1] .. [mk]\n\tsolve chinese reminder lemma\n");
    TRACE_THIS(" --crt-bn [k] [a1] .. [ak] [m1] .. [mk]\n\tsolve crt with BIGNUMs\n");
    TRACE_THIS(" --gen-keys <passphrase>\n\tgenerate DSA private/public key pair\n");
    TRACE_THIS(" --sha1 [filename]\n\tcalc SHA1 hash\n");
    TRACE_THIS(" --sign [filename] [prv_key pem] <passphrase>\n\tsign file\n");
    TRACE_THIS(" --sign-with-leak [filename] [prv_key pem] <passphrase>"
            "\n\tsign file, leaking 10 bits of private key through signature\n");
    TRACE_THIS(" --sign-with-msg [filename] [prv_key_pem] <passphrase> <sub message>"\
            "\n\tsign file, embedding a subliminal message in signature (k param)\n");
    TRACE_THIS(" --get-msg [filename] [filename.sig] [prv_key.pem] <passphrase>"\
            "\n\tget subliminal message from signature, usign private key\n");
    TRACE_THIS(" --verify [filename] [sig_file] [pub_key_pem]\n\tverify signature\n");
    TRACE_THIS("\nUtility functions\n");
    TRACE_THIS(" --enc-b64 [string]\n\tbase64 encode string\n");
    TRACE_THIS(" --dec-b64 [string]\n\tbase64 decode string\n");
    TRACE_THIS(" --is-quad-res [a] [p]\n\tis a quadratic residue modulo p? "\
            "(p positive, odd, prime)\n");
    TRACE_THIS("\ne.g.:   %s --exp-mod-bn 2 25306 25307\n", prog_name);
    TRACE_THIS("\t%s --crt 3 3 34 2 13 47 51 \n", prog_name);
    TRACE_THIS("\t%s --crt-bn 2 23 3232 551 7777 \n", prog_name);
}
